package com.ruoyi.ai.mongo;

import com.ruoyi.ai.mongo.pojo.ChatMessageDocument;
import com.ruoyi.ai.mongo.repository.ChatMessageRepository;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.time.LocalDateTime;
import java.util.List;


/**
 * MongoDB聊天记忆实现
 *
 * @author coach tam
 * @email 327395128@qq.com
 * @values 坚持灵活 灵活坚持
 * @since 2025/4/10
 */
@Slf4j
@Component
@Primary
public class MongoChatMemory implements ChatMemory {

    @Autowired
    private ChatMessageRepository chatMessageRepository;
    @Override
    public void add(String conversationId, List<Message> messages) {
        //将聊天记录保存到mongodb
        for (Message message : messages) {
            ChatMessageDocument chatMessageDocument = new ChatMessageDocument();
            chatMessageDocument.setUserId(conversationId);
            chatMessageDocument.setMessageType(message.getMessageType().getValue());
            chatMessageDocument.setText(message.getText());
            chatMessageDocument.setCreatedAt(LocalDateTime.now());
            chatMessageRepository.save(chatMessageDocument).subscribe();
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        log.info("从MongoDB获取对话记录, 对话ID: {}", conversationId);
        //查询条件为 userId=conversationId,根据createdAt降序排序
        Flux<ChatMessageDocument> flux = chatMessageRepository.findByUserIdOrderByCreatedAtDesc(conversationId);

        // 将flux转成List，避免在响应式线程中使用block
        try {
            return flux
                    .take(lastN)
                    .map(ChatMessageDocument::toMessage)
                    .collectList()
                    .toFuture()
                    .get();  // 使用CompletableFuture的get()替代block()
        } catch (Exception e) {
            log.error("获取聊天记录失败", e);
            return List.of();  // 出错时返回空列表
        }
    }

    @Override
    public void clear(String conversationId) {

    }
}
